import torch
import pandas as pd

import pandas as pd
import numpy as np

# This module is a CPU implement of reader.py

class Reader_CPU:
    # initial class Reader, load the needed SNP file
    def __init__(self):
        self.columns = []
        self.indexes = None

    # the CPU implement of readVCF method
    def readVCF(self, vcf_path, reset=False):
        # open the .vcf file
        with open(vcf_path, "r") as f:
            # if the line does not start as "##", this line is the header line
            header = next((i for i, line in enumerate(f) if not line.startswith('##')), 0)
        # read the file from header row
        df = pd.read_csv(vcf_path, header=header, sep='\t', dtype=str)
        # clean the useless columns, such as ['QUAL', 'FILTER', 'INFO', 'FORMAT', '#CHROM', 'POS', 'REF', 'ALT']
        # after doing that, only ID and sample IDs will be kept
        self.columns.append(df.columns[2])
        self.columns.extend(df.columns[9:])
        # concat the chromosome name and position as new ID for SNPs
        df['ID'] = df['#CHROM'].str.cat(df['POS'], sep='_')
        # set ID column as index and drop the useless columns
        df = df[self.columns].set_index('ID')
        # clear values of DataFrame
        # sometimes, values might be with annotations, for example, '1/1(some annotation)', which should be deleted
        for column in self.columns[1:]:
            # only the first three characters could be kept
            df[column] = df[column].str.slice(0, 3)
            # create a temporary series, which is replaced by rules
            temp_series = df[column].replace({r'1/1': '1', r'1|1': '1', r'0/1': '2', r'0|1': '2'})
            # make a mask to find out unchanged elements
            mask = df[column] == temp_series
            # assign '3' for unchanged elements
            temp_series[mask] = '3'
            # replace column of DataFrame with assigned series
            df[column] = temp_series.astype(np.int32)
            # delete useless variables
            del temp_series, mask
        # transpose the DataFrame, now the SNP IDs are columns, and sample IDs are indexes
        df = df.transpose()
        # keep the indexes and columns
        self.indexes = df.index
        self.columns = df.columns
        # drop indexes and reduce memory consumption
        if reset:
            df.reset_index(drop=True, inplace=True)
        return df.astype(np.int32)

columns = []
vcf_path = "data/train_example.vcf"
with open(vcf_path, "r") as f:
    # if the line does not start as "##", this line is the header line
    header = next((i for i, line in enumerate(f) if not line.startswith('##')), 0)

df = pd.read_csv(vcf_path, header=header, sep='\t', dtype=str)
columns.append(df.columns[2])
columns.extend(df.columns[9:])
# print(columns)
